package org.dark.web.scanner;

import com.baomidou.mybatisplus.annotation.TableName;
import com.baomidou.mybatisplus.core.mapper.Mapper;
import com.baomidou.mybatisplus.extension.toolkit.SqlHelper;
import org.apache.ibatis.session.SqlSession;
import org.dromara.streamquery.stream.core.clazz.ClassHelper;
import org.dromara.streamquery.stream.plugin.mybatisplus.Database;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TableScanner {

    final List<Class<?>> entityClassList;
    final Map<String, Class<?>> entityClassMap;



    public TableScanner(String basePackage) {
        entityClassList = ClassHelper.scanClasses(basePackage);
        entityClassMap = entityClassList.stream()
                .filter(clazz -> Optional.ofNullable(clazz.getAnnotation(TableName.class)).isPresent())
                .collect(Collectors.toMap(clazz -> (clazz.getAnnotation(TableName.class)).value(), Function.identity()));
    }

    public Mapper<?> getMapper(String entity) {
        Class<?> clazz = getEntity(entity);
        SqlSession sqlSession = SqlHelper.sqlSession(clazz);
        return Database.getMapper(clazz, sqlSession);
    }

    public Class<?> getEntity(String entity) {
        return entityClassMap.get(entity);
    }

}
